from tensorflow import keras
from .config import *
import tensorflow as tf

def build_model():
    policy = tf.keras.mixed_precision.Policy('mixed_float16')
    tf.keras.mixed_precision.set_global_policy(policy)
    input_image = keras.Input(shape=(config.IMG_SIDE_LEN, config.IMG_SIDE_LEN, config.NUM_COLOR_CHANNEL))
    input_pv_lstm = keras.Input(shape=(config.NUM_LOG_TERM, 1))
    input_pv_mlp =  keras.Input(shape=(config.NUM_LOG_TERM,))
    if config.MODEL_SELECT == "CNN_LSTM":
        if config.WEATHER_SELECT == 'CLOUDY':
            # CNN分支
            x = keras.layers.Conv2D(32, (3,3), activation='relu', padding='same')(input_image)
            x = keras.layers.MaxPooling2D(2)(x)
            x = keras.layers.Flatten()(x)
            x = keras.layers.Dropout(0.4)(x)
            cnn_output = keras.layers.Dense(32, activation='relu')(x)
            # LSTM分支 
            y = keras.layers.LSTM(64, return_sequences=True)(input_pv_lstm)
            y = keras.layers.Dropout(0.2)(y)
            y = keras.layers.LSTM(64, return_sequences=True)(y)
            y = keras.layers.Dropout(0.2)(y)
            y = keras.layers.LSTM(32, return_sequences=False)(y)
            y = keras.layers.Dropout(0.2)(y)
            lstm_output = keras.layers.Dense(32, activation='relu')(y)
            # 合并分支
            combined = keras.layers.concatenate([cnn_output, lstm_output])
            z = keras.layers.Dense(64, activation='relu')(combined)
            z = keras.layers.Dropout(0.2)(z)
            output = keras.layers.Dense(units=1)(z)
            model = keras.Model(inputs=[input_image, input_pv_lstm], outputs=output)
        elif config.WEATHER_SELECT == 'SUNNY':
            input_pv_lstm = keras.Input(shape=(config.NUM_LOG_TERM, 1))
            x = keras.layers.LSTM(64, return_sequences=False)(input_pv_lstm)
            x = keras.layers.Dropout(0.2)(x)
            x = keras.layers.Dense(16, activation='relu')(x)
            x = keras.layers.Dropout(0.2)(x)
            output = keras.layers.Dense(units=1)(x)
            model = keras.Model(inputs=input_pv_lstm, outputs=output)
        elif config.WEATHER_SELECT == 'OVERCAST':
            input_pv_lstm = keras.Input(shape=(config.NUM_LOG_TERM, 1))
            x = keras.layers.LSTM(32, return_sequences=False)(input_pv_lstm)
            x = keras.layers.Dropout(0.2)(x)
            x = keras.layers.Dense(16, activation='relu')(x)
            x = keras.layers.Dropout(0.2)(x)
            output = keras.layers.Dense(units=1)(x)
            model = keras.Model(inputs=input_pv_lstm, outputs=output)
        elif config.WEATHER_SELECT == 'ALL':
            input_pv_lstm = keras.Input(shape=(config.NUM_LOG_TERM, 1))
            x = keras.layers.LSTM(64, return_sequences=False)(input_pv_lstm)
            x = keras.layers.Dropout(0.2)(x)
            x = keras.layers.Dense(16, activation='relu')(x)
            x = keras.layers.Dropout(0.2)(x)
            output = keras.layers.Dense(units=1)(x)
            model = keras.Model(inputs=input_pv_lstm, outputs=output)
    elif config.MODEL_SELECT == "LSTM":
        if config.WEATHER_SELECT == 'ALL':
            input_pv_lstm = keras.Input(shape=(config.NUM_LOG_TERM, 1))
            x = keras.layers.LSTM(64, return_sequences=False)(input_pv_lstm)
            x = keras.layers.Dropout(0.2)(x)
            x = keras.layers.Dense(16, activation='relu')(x)
            x = keras.layers.Dropout(0.2)(x)
            output = keras.layers.Dense(units=1)(x)
            model = keras.Model(inputs=input_pv_lstm, outputs=output)
        elif config.WEATHER_SELECT == 'SUNNY':
            input_pv_lstm = keras.Input(shape=(config.NUM_LOG_TERM, 1))
            x = keras.layers.LSTM(64, return_sequences=False)(input_pv_lstm)
            x = keras.layers.Dropout(0.2)(x)
            x = keras.layers.Dense(16, activation='relu')(x)
            x = keras.layers.Dropout(0.2)(x)
            output = keras.layers.Dense(units=1)(x)
            model = keras.Model(inputs=input_pv_lstm, outputs=output)
        elif config.WEATHER_SELECT == 'CLOUDY':
            input_pv_lstm = keras.Input(shape=(config.NUM_LOG_TERM, 1))
            x = keras.layers.LSTM(64, return_sequences=False)(input_pv_lstm)
            x = keras.layers.Dropout(0.2)(x)
            x = keras.layers.Dense(16, activation='relu')(x)
            x = keras.layers.Dropout(0.2)(x)
            output = keras.layers.Dense(units=1)(x)
            model = keras.Model(inputs=input_pv_lstm, outputs=output)
        elif config.WEATHER_SELECT == 'OVERCAST':
            input_pv_lstm = keras.Input(shape=(config.NUM_LOG_TERM, 1))
            x = keras.layers.LSTM(32, return_sequences=False)(input_pv_lstm)
            x = keras.layers.Dropout(0.2)(x)
            x = keras.layers.Dense(16, activation='relu')(x)
            x = keras.layers.Dropout(0.2)(x)
            output = keras.layers.Dense(units=1)(x)
            model = keras.Model(inputs=input_pv_lstm, outputs=output)
    elif config.MODEL_SELECT == "MLP":
        # Fully connected layers
        x = keras.layers.Dense(64, activation='relu')(input_pv_mlp)
        x = keras.layers.Dropout(0.2)(x)
        x = keras.layers.Dense(32, activation='relu')(x)
        x = keras.layers.Dropout(0.2)(x)
        x = keras.layers.Dense(16, activation='relu')(x)
        x = keras.layers.Dropout(0.2)(x)
        # Regression to prediction target
        output = keras.layers.Dense(units=1)(x)
        # Construct the model
        model = keras.Model(inputs=input_pv_mlp, outputs=output)
    elif config.MODEL_SELECT == "U-Net":
        num_filters = 12
        # x_context: input context image frames
        x_in = keras.Input(shape=[64,64,3])
        
        # building block function of the generator
        #def ContractingBlock(X, input_channel):
        def Conv3x3block(X, input_channel):
            '''
            input_channel: number of input channels, a scaler
            '''
            conv1 = keras.layers.Conv2D(
                filters=input_channel,
                kernel_size=(3,3),
                padding="same")(X)
            bn1 = keras.layers.BatchNormalization()(conv1)
            relu1 = keras.layers.ReLU()(bn1)
            
            return relu1
        
        # building block function of the generator
        def BottleNeckBlock(X, input_channel):
            '''
            input_channel: number of input channels, a scaler
            '''
            conv1 = keras.layers.Conv2D(
                filters=input_channel,
                kernel_size=(3,3),
                padding="same")(X)
            bn1 = keras.layers.BatchNormalization()(conv1)
            relu1 = keras.layers.ReLU()(bn1)

            conv2 = keras.layers.Conv2D(    
                filters=input_channel,
                kernel_size=(3, 3),
                padding="same")(relu1)
            bn2 = keras.layers.BatchNormalization()(conv2)

            output = X + bn2

            return output
        
        # building block function of the generator
        #def ExpandingBlock(X, input_channel, skip_connection):
        def Up2x2Conv3x3(X, input_channel):
            '''
            input_channel: number of input channels, a scaler
            '''
            upsampling = keras.layers.UpSampling2D(size=(2, 2))(X)

            conv1 = keras.layers.Conv2D(
                filters=input_channel,
                kernel_size=(3, 3),
                padding="same")(upsampling)
            return conv1
            
        ## conv1x1 block convert 64*64*3 image to 64*64*num_filters image
        output1 = keras.layers.Conv2D(
            filters=num_filters,
            kernel_size=(1,1),
            padding="same")(x_in)

        ## Conv3x3block operation
        ### input shape: (64,64,num_filters); output shape: (64,64,num_filters)
        #output2 = Conv3x3block(output1, num_filters)
        output3 = Conv3x3block(output1, num_filters) # output3 will be skipped connected

        ## max pooling layer
        ### input shape: (64,64,num_filters); output shape: (32,32,num_filters)
        output4 = keras.layers.MaxPooling2D(pool_size=[2, 2], strides=2)(output3) 
        
        ## 2*Conv3x3block operation
        ### input shape: (32,32,num_filters); output shape: (32,32,2*num_filters)
        #output5 = Conv3x3block(output4, 2*num_filters)
        output6 = Conv3x3block(output4, 2*num_filters) # output6 will be skipped connected
        
        ## max pooling layer
        ### input shape: (32,32,2*num_filters); output shape: (16,16,2*num_filters)
        output7 = keras.layers.MaxPooling2D(pool_size=[2, 2], strides=2)(output6)
        
        ## Conv3x3block operation
        ### input shape: (32,32,2*num_filters); output shape: (16,16,4*num_filters)
        output8 = Conv3x3block(output7, 4*num_filters)
        
        ## Bottleneck - Residual blocks
        ### input shape: (16,16,4*num_filters); output shape: (16,16,4*num_filters)
        #output8 = BottleNeckBlock(output7, 4*num_filters)
        output9 = BottleNeckBlock(output8, 4*num_filters)
        output10 = BottleNeckBlock(output9, 4*num_filters)

        ## Up2x2Conv3x3 block
        ### input shape: (16,16,4*num_filters); output shape: (32,32,2*num_filters)
        output11 = Up2x2Conv3x3(output10, 2*num_filters)
        
        ## Concatenation
        ### input shape: (32,32,2*num_filters); output shape: (32,32,4*num_filters)
        output12 = keras.layers.concatenate([output11, output6], axis=3)
        
        ## 2*Conv3x3block operation
        ### input shape: (32,32,4*num_filters); output shape: (32,32,2*num_filters)
        #output13 = Conv3x3block(output12, 2*num_filters)
        #output13 = keras.layers.Dropout(0.4)(output13)
        output14 = Conv3x3block(output12, 2*num_filters)
        output14 = keras.layers.Dropout(0.4)(output14)
        
        ## Up2x2Conv3x3 block
        ### input shape: (32,32,2*num_filters); output shape: (64,64,num_filters)
        output15 = Up2x2Conv3x3(output14,num_filters)
        
        ## Concatenation
        ### input shape: (64,64,num_filters); output shape: (64,64,2*num_filters)
        output16 = keras.layers.concatenate([output15, output3], axis=3)
        
        ## 2*Conv3x3block operation
        ### input shape: (64,64,2*num_filters); output shape: (64,64,num_filters)
        #output17 = Conv3x3block(output16, num_filters)
        #output17 = keras.layers.Dropout(0.4)(output17)
        output18 = Conv3x3block(output16, num_filters)
        output18 = keras.layers.Dropout(0.4)(output18)

        ## Conv1x1 block convert 64*64*12 image to 64*64*1 image
        ### y_out
        y = keras.layers.Conv2D(
            filters=1,
            kernel_size=(1,1),
            padding="same",
            activation='relu')(output18)
        
        unet_output = keras.layers.Flatten()(y)
        y = keras.layers.LSTM(64, return_sequences=True)(input_pv_lstm)
        y = keras.layers.Dropout(0.2)(y)
        y = keras.layers.LSTM(32, return_sequences=False)(y)
        y = keras.layers.Dropout(0.2)(y)
        lstm_output = keras.layers.Dense(32, activation='relu')(y)
        # 合并分支
        combined = keras.layers.concatenate([unet_output, lstm_output])
        z = keras.layers.Dense(64, activation='relu')(combined)
        z = keras.layers.Dropout(0.2)(z)
        output = keras.layers.Dense(units=1)(z)
        
        # construct the model
        model = keras.Model(inputs=[x_in,input_pv_lstm],outputs=output,name="Image2PV")
    else:
        print("Selectd model is error.")
    return model 

